# -*- coding: utf-8 -*-
import sys
import os
import time

def train_model():
    try:
        # 导入ultralytics
        try:
            from ultralytics import YOLO
            print("成功导入ultralytics")
        except ImportError as e:
            print(f"导入ultralytics失败: {e}")
            print("请安装: pip install ultralytics")
            return
            
        # 加载模型
        print("加载模型: D:/VsCodeProject/ultralytics02/yolov8n.pt")
        try:
            model = YOLO('D:/VsCodeProject/ultralytics02/yolov8n.pt')
            print("模型加载成功")
        except Exception as e:
            print(f"模型加载失败: {e}")
            return
        
        # 训练配置
        train_config = {'agnostic_nms': False, 'amp': True, 'augment': False, 'auto_augment': 'randaugment', 'batch': 8, 'bgr': 0, 'box': 7.5, 'cache': False, 'cfg': None, 'classes': None, 'close_mosaic': 10, 'cls': 0.5, 'conf': None, 'copy_paste': 0, 'copy_paste_mode': 'flip', 'cos_lr': False, 'cutmix': 0, 'data': 'D:\\VsCodeProject\\ultralytics02\\fire_warning\\data.yaml', 'degrees': 0, 'deterministic': True, 'device': None, 'dfl': 1.5, 'dnn': False, 'dropout': 0, 'dynamic': False, 'embed': None, 'epochs': 100, 'erasing': 0.4, 'exist_ok': False, 'fliplr': 0.5, 'flipud': 0, 'format': 'torchscript', 'fraction': 1, 'freeze': None, 'half': False, 'hsv_h': 0.015, 'hsv_s': 0.7, 'hsv_v': 0.4, 'imgsz': 640, 'int8': False, 'iou': 0.7, 'keras': False, 'kobj': 1, 'line_width': None, 'lr0': 0.01, 'lrf': 0.01, 'mask_ratio': 4, 'max_det': 300, 'mixup': 0, 'mode': 'train', 'model': 'yolov8n.pt', 'momentum': 0.937, 'mosaic': 1, 'multi_scale': False, 'name': None, 'nbs': 64, 'nms': False, 'opset': None, 'optimize': False, 'optimizer': 'auto', 'overlap_mask': True, 'patience': 100, 'perspective': 0, 'plots': True, 'pose': 12, 'pretrained': True, 'profile': False, 'project': None, 'rect': False, 'resume': False, 'retina_masks': False, 'save': True, 'save_conf': False, 'save_crop': False, 'save_frames': False, 'save_json': False, 'save_period': -1, 'save_txt': False, 'scale': 0.5, 'seed': 0, 'shear': 0, 'show': False, 'show_boxes': True, 'show_conf': True, 'show_labels': True, 'simplify': True, 'single_cls': False, 'source': None, 'split': 'val', 'stream_buffer': False, 'task': 'detect', 'time': None, 'tracker': 'botsort.yaml', 'translate': 0.1, 'val': True, 'verbose': True, 'vid_stride': 1, 'visualize': False, 'warmup_bias_lr': 0.1, 'warmup_epochs': 3, 'warmup_momentum': 0.8, 'weight_decay': 0.0005, 'workers': 8, 'workspace': None}
        train_config['data'] = r'D:/VsCodeProject/ultralytics02/fire_warning/data.yaml'
        train_config['project'] = r'D:/VsCodeProject/ultralytics02/fire_warning/runs'
        train_config['name'] = 'yolo_training'
        
        print("训练配置:")
        for key, value in train_config.items():
            print(f"  {key}: {value}")
        
        print("开始训练...")
        # 开始训练
        results = model.train(**train_config)
        
        print("训练完成!")
        print("结果保存在:", results)
        
    except Exception as e:
        print(f"训练出错: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    train_model()
